use std::any::{Any, TypeId};
use dashmap::DashMap;
use crate::message::{Message, MessageHeader};
use crate::errors::SerializationError;

// Type-erased serialization trait
pub trait DynSerializer: Send + Sync {
    fn serialize_dyn(&self, message: &dyn Any) -> Result<Vec<u8>, SerializationError>;
    fn deserialize_dyn(&self, bytes: &[u8], type_id: TypeId) -> Result<Box<dyn Any>, SerializationError>;
}

// Generic serializer trait
pub trait Serializer<T: Any + Send + Clone>: Send + Sync {
    fn serialize(&self, message: &Message<T>) -> Result<Vec<u8>, SerializationError>;
    fn deserialize(&self, bytes: &[u8]) -> Result<Message<T>, SerializationError>;
}

// Adapter to convert generic Serializer to DynSerializer
pub struct SerializerAdapter<T: Any + Send + Clone, S: Serializer<T>> {
    serializer: S,
    _phantom: std::marker::PhantomData<T>,
}

impl<T: Any + Send + Clone, S: Serializer<T>> SerializerAdapter<T, S> {
    pub fn new(serializer: S) -> Self {
        Self {
            serializer,
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<T: Any + Send + Clone, S: Serializer<T>> DynSerializer for SerializerAdapter<T, S> {
    fn serialize_dyn(&self, message: &dyn Any) -> Result<Vec<u8>, SerializationError> {
        if let Some(msg) = message.downcast_ref::<Message<T>>() {
            self.serializer.serialize(msg)
        } else {
            Err(SerializationError::TypeMismatch)
        }
    }

    fn deserialize_dyn(&self, bytes: &[u8], type_id: TypeId) -> Result<Box<dyn Any>, SerializationError> {
        if type_id == TypeId::of::<Message<T>>() {
            let msg = self.serializer.deserialize(bytes)?;
            Ok(Box::new(msg))
        } else {
            Err(SerializationError::TypeMismatch)
        }
    }
}

pub struct SerializerRegistry {
    serializers: DashMap<TypeId, Box<dyn DynSerializer>>,
}

impl SerializerRegistry {
    pub fn new() -> Self {
        Self {
            serializers: DashMap::new(),
        }
    }

    pub fn register<T: Any + Send + Clone, S: Serializer<T> + 'static>(&self, serializer: S) {
        let adapter = SerializerAdapter::new(serializer);
        self.serializers.insert(TypeId::of::<T>(), Box::new(adapter));
    }

    pub fn serialize<T: Any + Send + Clone>(&self, message: &Message<T>) -> Result<Vec<u8>, SerializationError> {
        if let Some(serializer) = self.serializers.get(&TypeId::of::<T>()) {
            serializer.serialize_dyn(message)
        } else {
            Err(SerializationError::NoSerializerFound)
        }
    }

    pub fn deserialize<T: Any + Send + Clone>(&self, bytes: &[u8]) -> Result<Message<T>, SerializationError> {
        if let Some(serializer) = self.serializers.get(&TypeId::of::<T>()) {
            let dyn_msg = serializer.deserialize_dyn(bytes, TypeId::of::<Message<T>>())?;
            dyn_msg.downcast::<Message<T>>()
                .map(|b| *b)
                .map_err(|_| SerializationError::TypeMismatch)
        } else {
            Err(SerializationError::NoSerializerFound)
        }
    }
} 